// stdlib                  http://www.tomshardware.com/
#include <cstdlib>
#include <vector>
#include <string>
#include <fstream>
#include <iostream>

// boost
#include <boost/format.hpp>
#include <boost/lexical_cast.hpp>

// components
#include <tree.h>
#include <print_error.h>
#include <tree_ops.h>
#include <portable_timer.h>
#include <perlutil.h>
#define _MY_STL_MACROS
#include <stlhelper.h>
#include <open_stream_name_with_substitutions.h>
#define DEBUG_LEVEL 0
#include <debug_func.h>
#include <progress_indicator.h>


#include "get_arg.h"

using boost::format;
using boost::lexical_cast;
using std::vector;
using std::string;
using std::ofstream;
using std::cout;
using std::cin;
using std::cerr;
using std::replace;

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	print and analyse orthologs

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
void open_filename(ofstream& ofs, const string& path_format, const string& base_file_name)
{
	debug_func_cerr( 5 );
	std::string new_file_name;
	open_stream_with_substitutions(	ofs, new_file_name, path_format, base_file_name);
}

struct t_ortholog_set_gene_counts
{
	unsigned gene_set_size;
	unsigned cnt_tax1;
	unsigned cnt_tax2;
	t_ortholog_set_gene_counts():gene_set_size(0), cnt_tax1(0), cnt_tax2(0){}
};



void print_genes_from_ortholog_set(ofstream& ofs, const string& path_format, const string& file_name,
							const string& file_suffix,
							const string& header, const bcpp_tree_ops::t_ortholog_data& data)
{
	debug_func_cerr( 5 );
	if (!data.gene_sets.size())
		return;

	if (!ofs.is_open())
		open_filename(ofs, path_format, file_name + file_suffix);


	ofs << header << std::endl;
	if (data.gene_sets[0].length > 0.0)
	{
		// orthology nodes, length saved
		for (unsigned i = 0; i < data.gene_sets.size(); ++i)
		{
			ofs << data.gene_sets[i].length;
			for (unsigned j = 0 ; j < data.gene_sets[i].genes.size(); ++j)
				ofs << "\t" << data.gene_sets[i].genes[j];
			ofs << "\n";
		}
	}
	else
	{
		// non-orthology nodes, no length saved
		for (unsigned i = 0; i < data.gene_sets.size(); ++i)
		{
			if (data.gene_sets[i].genes.size())
				ofs << data.gene_sets[i].genes[0];
			for (unsigned j = 1 ; j < data.gene_sets[i].genes.size(); ++j)
				ofs << "\t" << data.gene_sets[i].genes[j];
			ofs << "\n";
		}
	}
	ofs << std::endl;
	return;
}


void print_orphaned_clades(ostream& ofs, const bcpp_tree_ops::t_ortholog_data& data)
{
	debug_func_cerr( 5 );
	if (!data.gene_sets.size())
		return;

	// orphaned nodes, no length saved
	for (unsigned i = 0; i < data.gene_sets.size(); ++i)
	{
		if (!data.gene_sets[i].genes.size())
			continue;
		ofs << data.gene_sets[i].genes[0];
		for (unsigned j = 1 ; j < data.gene_sets[i].genes.size(); ++j)
			ofs << "\t" << data.gene_sets[i].genes[j];
		ofs << std::endl;
	}
	return;
}


void print_all_genes_from_ortholog_set(ostream& ofs, const bcpp_tree_ops::t_ortholog_data& data)
{
	debug_func_cerr( 5 );
	if (!data.gene_sets.size())
		return;

	// orthology nodes, length saved
	for (unsigned i = 0; i < data.gene_sets.size(); ++i)
	{
		ofs << data.gene_sets[i].length;
		for (unsigned j = 0 ; j < data.gene_sets[i].genes.size(); ++j)
			ofs << "\t" << data.gene_sets[i].genes[j];
		ofs << std::endl;
	}
	return;
}


struct t_outputstreams
{

	t_outputstreams(const string& prefix, const string& orphaned_clades_) : 
			path_format(prefix),
			orphaned_clades(orphaned_clades_)
	{
		debug_func_cerr( 5 );
		open();

		ofs_all_nodes	<< "To root\tTo leaf\tGene id\tNode type\n"
						<< string(80, '-') << "\n";

		ofs_full_tree	<< "[Branch length / To root / To leaf] Gene id {Node type}\n"
						<< string(80, '-') << "\n";
	}
	string	 path_format;
	string	 orphaned_clades;
	ofstream ofs_ortholog_sets[bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT];
//	ofstream ofs_all_orthologs;
	ofstream ofs_full_tree;
	ofstream ofs_all_nodes;
	ofstream ofs_orphaned_clades;
	void open();
	bool printout_analysis(	const bcpp_tree::t_tree&			tree,
							const string&						tax1,
							const string&						tax2,
							vector<t_ortholog_set_gene_counts>& gene_counts,
							bool								breakdown_by_ortholog_type);
};


bool t_outputstreams::printout_analysis(const bcpp_tree::t_tree& tree,
										const string& tax1, const string& tax2,
										vector<t_ortholog_set_gene_counts>& gene_counts,
										bool breakdown_by_ortholog_type)
{
	debug_func_cerr( 5 );
	string header  = ">" + tree.get_name();
	bcpp_tree_ops::do_calculate_branch_lengths tree_branch_lengths(tree);

	// set type for each node
	bcpp_tree_ops::do_set_node_types tree_orthology(tree, tax1, tax2);

	// dump all nodes
	{
		ofs_all_nodes << header << std::endl;
		bcpp_tree_ops::op_print_tree dump_all_nodes
											(tree,
											ofs_all_nodes,
											tree_branch_lengths.lengths_to_leaves,
											tree_branch_lengths.lengths_to_root,
											tree_orthology.node_types,
											true // nodes only
											);
		ofs_all_nodes << std::endl;
	}

	// print full tree
	{

		ofs_full_tree << header << std::endl;
		bcpp_tree_ops::op_print_tree_elaborate print_tree_elaborate
											(tree,
											ofs_full_tree,
											tree_branch_lengths.lengths_to_leaves,
											tree_branch_lengths.lengths_to_root,
											tree_orthology.node_types,
											3	// precision
											);
		print_tree_elaborate.print ();
		ofs_full_tree << std::endl;
	}



	// print gene sets
	{
		bcpp_tree_ops::do_collect_orthologs data(tree_orthology.node_types,
													tree_branch_lengths.lengths_to_leaves);
		tree.for_each(data);

		//
		//	print lists of genes by orthologue type
		//
		if (!breakdown_by_ortholog_type)
		{
			std::string file_names[bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT] =
												{
														"ortho_1_to_1"    ,
														"ortho_1_to_m"    ,
														"ortho_m_to_1"    ,
														"ortho_m_to_m"    ,
														"orphaned_tax1"   ,
														"orphaned_tax2"  };
	
			for (unsigned i = 0; i < bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT; ++i)
			{
				print_genes_from_ortholog_set(ofs_ortholog_sets[i], path_format,
												   file_names[i], ".data",
													header, data.ortholog_sets[i]);
			}
		}

		// print out ortholog clades to cout, and orphaned clades to specified file
		for (unsigned i = bcpp_tree_ops::do_collect_orthologs::e_ORTH_1_TO_1; 
			  i <= bcpp_tree_ops::do_collect_orthologs::e_ORTH_M_TO_M; ++i)
		{	  
			print_all_genes_from_ortholog_set(cout, data.ortholog_sets[i]);
		}									   
		for (unsigned i = bcpp_tree_ops::do_collect_orthologs::e_ORPH_TAX1; 
			  i <= bcpp_tree_ops::do_collect_orthologs::e_ORPH_TAX2; ++i)
		{
			print_orphaned_clades(ofs_orphaned_clades, data.ortholog_sets[i]);
		}
		
		for (unsigned i = 0; i < bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT; ++i)
		{
			gene_counts[i].cnt_tax1 += data.ortholog_sets[i].cnt_tax1;
			gene_counts[i].cnt_tax2 += data.ortholog_sets[i].cnt_tax2;
			gene_counts[i].gene_set_size += data.ortholog_sets[i].gene_sets.size();
		}

	}
	return true;

	/*   ofs_full_tree << header << std::endl;
	   bcpp_tree_ops::op_print_tree print_tree
										   (tree,
										   ofs_full_tree,
										   tree_branch_lengths.lengths_to_leaves,
										   tree_branch_lengths.lengths_to_root,
										   tree_orthology.node_types,
										   false // nodes only
										   );
	   ofs_full_tree << std::endl;
	  */
	   /*
	   // print new_hampshire tree
	   {
		   ofs_newhampshire_tree << header << std::endl;
		   bcpp_tree_ops::do_newhampshire_format_str nh_str(ofs_newhampshire_tree);
		   tree.for_each(nh_str);
		   ofs_newhampshire_tree << std::endl;
	   }
	   */


}



void t_outputstreams::open()
{
	debug_func_cerr( 5 );
	open_filename(ofs_full_tree,   path_format, "full_topology.tree" );
	open_filename(ofs_all_nodes,   path_format, "all_nodes.list"	 );

	if (orphaned_clades.length())
	{
		ofs_orphaned_clades.open(orphaned_clades.c_str());
		if (!ofs_orphaned_clades)
			throw std::runtime_error("Could not open \"" + orphaned_clades + "\"");
	}
}

//________________________________________________________________________________________

// print_gene_counts

//________________________________________________________________________________________
void print_gene_counts(vector<t_ortholog_set_gene_counts>& gene_counts, t_program_args& args)
{
	ofstream ofs_gene_counts;
	std::string file_name;
	open_stream_with_substitutions(	ofs_gene_counts, file_name, args.out_path_format, "ortholog.counts");
	args.VERBOSE() << "\tPrinting ortholog counts to \"" << file_name << "\"\n";
	std::string ortholog_type_names[bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT] =
										{
												string("1 to 1 orthologs")   	 ,
												string("1 to m orthologs")   	 ,
												string("m to 1 orthologs")   	 ,
												string("m to m orthologs")   	 ,
												args.taxon1_name + " orphans"	 ,
												args.taxon2_name + " orphans"};
	// make all categories the same length
	unsigned max_len = 0;
	for (unsigned i = 0; i < bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT; ++i)
		max_len = std::max(max_len, ortholog_type_names [i].length());
	for (unsigned i = 0; i < bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT; ++i)
		if (max_len !=ortholog_type_names[i].size())
			ortholog_type_names[i].append(string(max_len - ortholog_type_names[i].length(), ' '));

	ofs_gene_counts	<< string(max_len, ' ') + "\t"
					<< "Ortholog Sets\t"
					<< args.taxon1_name << "\t"
					<< args.taxon2_name << "\n"
					<< string(max_len, ' ') + "\t"
					<< "--------\t--------\t--------" << std::endl;
	unsigned total_gene_cnt_tax1 = 0;
	unsigned total_gene_cnt_tax2 = 0;
	unsigned total_gene_sets = 0;
	for (unsigned i = 0; i < bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT; ++i)
	{
		if (gene_counts[i].gene_set_size)
		{
			ofs_gene_counts	<< ortholog_type_names[i] << "\t"
							<< gene_counts[i].gene_set_size << "\t"
							<< gene_counts[i].cnt_tax1 << "\t"
							<< gene_counts[i].cnt_tax2 << std::endl;
			total_gene_sets += gene_counts[i].gene_set_size;
			total_gene_cnt_tax1 += gene_counts[i].cnt_tax1;
			total_gene_cnt_tax2 += gene_counts[i].cnt_tax2;
		}
	}
	ofs_gene_counts
						<< string(max_len, ' ') + "\t"
						<< "--------\t--------\t--------\n"
						<< "Total" + string(max_len - 5, ' ') + "\t"
						<< total_gene_sets
						<< "\t" << total_gene_cnt_tax1
						<< "\t" << total_gene_cnt_tax2 << std::endl;

}

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
int main (int argc, char *argv[])
{
	debug_func_cerr( 5 );
	// we use cout to print progress so flush per call
	//cout.setf(std::ios_base::unitbuf);

	std::ios::sync_with_stdio(false);

	t_program_args args;
	try
	{
	if (!process_command_line_options(argc, argv, args))
		return 1;



	t_outputstreams strms (args.out_path_format, args.orphaned_clades);

	unsigned cnt_tree = 0;
	unsigned cnt_lines = 0;
	vector<t_ortholog_set_gene_counts> gene_counts(bcpp_tree_ops::do_collect_orthologs::eORTH_TYPE_CNT);


	args.VERBOSE() << "\tDerive orthology from trees\n";
	t_progress_indicator dots(args.VERBOSE(1) , 2000);
	while (cin)
	{
		bcpp_tree::t_tree paralog_tree;
		cin >> bcpp_tree::t_tree_io_wrapper(paralog_tree, cnt_lines);
		if (!cin)
			break;
		++cnt_tree;
		if (paralog_tree.get_name().empty())
			paralog_tree.set_name("tree #" + ultoa(cnt_tree));

		// calculate orthology and print
		if (!strms.printout_analysis(paralog_tree, args.taxon1, args.taxon2, gene_counts, 
									 !args.no_breakdown_by_ortholog_type))
				return 0;
		++dots;
	}
	dots.finish();


	print_gene_counts(gene_counts, args);

	args.VERBOSE() << "\tAll finished!!\n\n";
	}
	catch(std::exception& e)
	{
		std_print_error(args.CERR(), e.what());
		return 1;
	}

	return 0;
}

